import os.path
import random
import sys

import numpy as np
import pygame

WIDTH, HEIGHT = 488, 640  # 窗口的宽高
NUMICON = 270  # 图标数量  必须是3的倍数 图标种类的倍数（animal种类=15 fruit种类=10）
SIZE = 48  # 图标大小
MINX = 76  # 图标最左位置
MAXX = WIDTH - SIZE - MINX  # 图标最右位置
MINY = MINX  # 图标最上位置
MAXY = MINY + 6 * SIZE  # 图标最下位置
SELECTIONY = MAXY + SIZE * 2
TOOLY = SELECTIONY + SIZE * 2


def draw_bbox(screen, rect, color=(0, 0, 0), width=2):
    """ 绘制指定区域的边界框"""
    pygame.draw.rect(screen, color, rect, width)


def draw_line(screen, start, end, color=(0, 0, 0), width=1):
    """ 绘制指定线段 """
    pygame.draw.line(screen, color, start, end, width)


def box_iou_xywh(box1, box2):
    """ 计算两个矩形框的交并比 用于判断两个矩形框是否有重合部分"""
    x1min, y1min = box1[0], box1[1]
    x1max, y1max = box1[0] + box1[2], box1[1] + box1[3]
    s1 = box1[2] * box1[3]

    x2min, y2min = box2[0], box2[1]
    x2max, y2max = box2[0] + box2[2], box2[1] + box2[3]
    s2 = box2[2] * box2[3]

    xmin = np.maximum(x1min, x2min)
    ymin = np.maximum(y1min, y2min)
    xmax = np.minimum(x1max, x2max)
    ymax = np.minimum(y1max, y2max)
    inter_h = np.maximum(ymax - ymin, 0.)
    inter_w = np.maximum(xmax - xmin, 0.)
    intersection = inter_h * inter_w

    union = s1 + s2 - intersection
    iou = intersection / union
    return iou


class Icon(pygame.sprite.Sprite):
    """ 图标 """

    def __init__(self, idx, imgpath, position, level, over):
        """
        游戏画面中需要用到的图标
        :param idx: 图标的唯一索引
        :param imgpath: 图标的图片路径
        :param position: 图标的摆放位置 取左上角的点 即left, top
        :param level: 图标的摆放层级 用于模型的训练
        :param over:
        """
        pygame.sprite.Sprite.__init__(self)
        self.idx = idx  # 唯一标识一个图标
        self.imgpath = imgpath  # 该图标的图片路径
        self.image = pygame.image.load(imgpath)  # 加载图片
        self.image = pygame.transform.smoothscale(self.image, (SIZE, SIZE))  # 调整图片尺寸
        self.rect = self.image.get_rect()  # 位置信息
        self.rect.left, self.rect.top = position
        self.level = level  # 存储的层数 从1开始 逐渐往上堆叠 数字增加 （为了模型方便 后续反转层数由上往下依次增加层数）
        self.type = os.path.splitext(os.path.basename(imgpath))[0]  # 图标所属类别
        self.over = over  # 该图标覆盖的 非必要属性
        self.under = []  # 覆盖该图标的


def get_tool_by_position(position):
    """ 根据鼠标点击的位置获取道具 如果已使用道具 则点击无效 点击后更新为已使用道具"""
    x, y = position
    tool_x1 = MINX
    tool_x2 = MINX + SIZE * 3
    tool_x3 = MINX + SIZE * 6
    if TOOLY < y < TOOLY + SIZE:
        if tool_x1 < x < tool_x1 + SIZE:
            return 'put'
        if tool_x2 < x < tool_x2 + SIZE:
            return 'shuffle'
        if tool_x3 < x < tool_x3 + SIZE:
            return 'withdraw'
    return None


class Game:
    def __init__(self):
        self.window = None
        pygame.init()
        pygame.display.set_caption('爱消除')
        # 加载背景音乐
        pygame.mixer.init()
        pygame.mixer.music.load(os.path.join("resources/audios/bg.mp3"))
        pygame.mixer.music.set_volume(0.6)

        # 加载音效
        self.click = pygame.mixer.Sound(os.path.join('resources/audios/click.wav'))  # 点击图标时 播放音效
        self.match = pygame.mixer.Sound(os.path.join('resources/audios/match.wav'))  # 候选框中有匹配时 播放音效
        # 加载字体
        self.font = pygame.font.Font(pygame.font.get_default_font(), 18)
        # 加载图标
        self.imgpaths = []  # 所有图标元素的存储路径
        for icon in os.listdir('resources/icons/animal/bbox'):
            icon_path = os.path.join('resources/icons/animal/bbox', icon)
            self.imgpaths.append(icon_path)
        assert NUMICON % (3 * len(self.imgpaths)) == 0, \
            f'NUMICON={NUMICON}, can not be divided by 3 and the amount of classes.'
        # 图标横轴限定范围
        self.xrange = np.arange(MINX, MAXX + SIZE / 2, SIZE / 2)
        # 图标纵轴限定范围
        self.yrange = np.arange(MINY, MAXY + SIZE / 2, SIZE / 2)
        self.remaining = NUMICON  # 剩余图标数量
        self.maxlevel = 0  # 存储最大层数 用于将层数反转 不依赖层数来判断是否可点击 是否加阴影
        # 初始化 sprite group
        self.group = self._init_group()  # 未点击的图标
        self.selected = []  # 已点击的图标
        self.last_selected_icon_idx = None  # 最后选择的图标
        # 是否已使用道具 放回一张 取出三张 随机打乱
        self.used_put, self.used_withdraw, self.used_shuffle = False, False, False

        self.keep_use_tool = False  # 无限使用道具
        self.finished = False  # 游戏结束
        self.failed_game = False  # 游戏失败
        self.success_game = False  # 游戏成功

    def _click_sound(self):
        self.click.play()

    def _match_sound(self):
        self.match.play()

    def _draw_selection(self, pos):
        """绘制7个选择框"""
        left, top = pos
        # 绘制大矩形
        draw_bbox(self.window, (left, top, 7 * SIZE, SIZE))
        # 绘制虚线
        for i in range(6):
            x, y1, y2 = left + (i + 1) * SIZE, top, top + SIZE
            for j in np.linspace(y1, y2, 5, endpoint=False):
                draw_line(self.window, (x, j), (x, j + 4.8))

    def _draw_remaining_icon_amount(self):
        """ 绘制剩余的图标数量 """
        remaining = self.font.render(f'{self.remaining}', True, (198, 77, 63))
        rect = remaining.get_rect()
        rect.left, rect.top = (6, 6)
        self.window.blit(remaining, rect)

    def _draw_result(self):
        """ 绘制游戏结果 """
        result_text = ' '
        if self.success_game:
            result_text = 'SUCCESS!'
        if self.failed_game:
            result_text = 'Failed!'
        result = self.font.render(f'{result_text} press <r> key to restart.', True, (255, 255, 255))
        rect = result.get_rect(center=(WIDTH / 2, HEIGHT / 2))

        s = pygame.Surface((WIDTH, HEIGHT))  # the size of your rect
        s.set_alpha(200)  # alpha level
        s.fill((46, 46, 46,))  # this fills the entire surface
        s.blit(result, rect)
        self.window.blit(s, (0, 0))

    def _match(self):
        """ 判断候选框中的图标是否有匹配的 有则删除 """
        last_type, cnt, order_idx = 0, 0, -1  # 前一个图标的类别，前一个图标的计数，新一轮计数中的第一个图标索引
        for i in range(len(self.selected)):
            tmp_icon = self.selected[i]
            if tmp_icon.type == last_type:
                cnt += 1
                if cnt == 3:
                    # 有三个相同类别的图标 需要删除这三个图标
                    self._match_sound()
                    del self.selected[order_idx]  # 删除的是连续位置 所以每次删除一个元素之后 索引都会-1 故删除之后的元素对应索引不变
                    del self.selected[order_idx]
                    del self.selected[order_idx]
                    return  # 一次最多有三个相同类别需要检查 当执行到这里时 可以肯定已筛选完毕 去除return会索引越界 也可以用break
            else:
                # 开始新一轮计数
                last_type = tmp_icon.type
                cnt = 1
                order_idx = i

    def _draw_tools(self):
        """ 绘制道具 """
        tool_put = pygame.image.load(os.path.join('resources/tools/put.png'))  # 放回一张图标
        tool_put = pygame.transform.smoothscale(tool_put, (SIZE, SIZE))
        tool_shuffle = pygame.image.load(os.path.join('resources/tools/shuffle.png'))  # 随机打乱
        tool_shuffle = pygame.transform.smoothscale(tool_shuffle, (SIZE, SIZE))
        tool_withdraw = pygame.image.load(os.path.join('resources/tools/withdraw.png'))  # 取出三张图标
        tool_withdraw = pygame.transform.smoothscale(tool_withdraw, (SIZE, SIZE))

        self.window.blit(tool_put, (MINX, TOOLY))
        self.window.blit(tool_shuffle, (MINX + SIZE * 3, TOOLY))
        self.window.blit(tool_withdraw, (MINX + SIZE * 6, TOOLY))

        # 如果道具已经用过了 则添加阴影
        if self.used_put and not self.keep_use_tool:
            self._draw_transparency_block((MINX, TOOLY))
        if self.used_shuffle and not self.keep_use_tool:
            self._draw_transparency_block((MINX + SIZE * 3, TOOLY))
        if self.used_withdraw and not self.keep_use_tool:
            self._draw_transparency_block((MINX + SIZE * 6, TOOLY))

    def _draw_transparency_block(self, position):
        """ 给不能点击的图标加上半透明阴影 """
        s = pygame.Surface((SIZE, SIZE))  # the size of your rect
        s.set_alpha(128)  # alpha level
        s.fill((46, 46, 46))  # this fills the entire surface
        self.window.blit(s, position)

    def _init_group(self, num_of_icon=NUMICON):
        """ 初始化图标信息 随即设定位置和图片 根据位置 判断层数 初始化over和under """
        group = []
        for i in range(num_of_icon):
            idx = i
            imgpath = self.imgpaths[i % len(self.imgpaths)]
            left, top = random.choice(self.xrange), random.choice(self.yrange),
            # 判断层数：遍历已有的图标 判断是否有重合 没有则层数为1 有则＋1
            tmp_level = 0  # 临时存放已有的覆盖中 层数最高的那一个 （因为可能覆盖了多个图标 只存放层数最大的那个）
            over = []  # 覆盖的图标
            for tmp_icon in group:
                if box_iou_xywh((left, top, SIZE, SIZE), tmp_icon.rect) > (1 / 8):
                    # 有重合
                    tmp_icon.under.append(idx)
                    over.append(tmp_icon.idx)
                    if tmp_icon.level > tmp_level:
                        tmp_level = tmp_icon.level

            level = tmp_level + 1
            if level > self.maxlevel:
                self.maxlevel = level
            icon = Icon(idx, imgpath, (left, top), level, over)
            group.append(icon)
        return group

    def reset(self):
        """ 重置游戏"""
        self.remaining = NUMICON  # 重置剩余图标数量
        self.maxlevel = 0  # 重置最大层数
        self.group = self._init_group()  # 重置图标
        self.selected.clear()  # 清空已选择图标
        self.last_selected_icon_idx = None
        self.used_put = False
        self.used_withdraw = False
        self.used_shuffle = False
        self.keep_use_tool = False
        self.finished = False
        self.success_game = False
        self.failed_game = False

    def reverse_level(self):
        tmp_level = self.maxlevel + 1
        for icon in self.group:
            icon.level = tmp_level - icon.level
        for icon in self.selected:
            icon.level = tmp_level - icon.level

    def tool_shuffle(self):
        """ 打乱图标 """
        if not self.used_shuffle:
            # 储存已有图标的类别
            tmp_remaining_type = []
            for tmp_icon in self.group:
                tmp_remaining_type.append(tmp_icon.type)
            # 开始重新生成图标以及位置关系
            group = []
            for i in range(len(tmp_remaining_type)):
                idx = i
                imgpath = os.path.join('resources/icons/animal/bbox', tmp_remaining_type[i] + '.png')
                left, top = random.choice(self.xrange), random.choice(self.yrange),
                # 判断层数：遍历已有的图标 判断是否有重合 没有则层数为1 有则＋1
                tmp_level = 0  # 临时存放已有的覆盖中 层数最高的那一个 （因为可能覆盖了多个图标 只存放层数最大的那个）
                over = []  # 覆盖的图标
                for tmp_icon in group:
                    if box_iou_xywh((left, top, SIZE, SIZE), tmp_icon.rect) > (1 / 8):
                        # 有重合
                        tmp_icon.under.append(idx)
                        over.append(tmp_icon.idx)
                        if tmp_icon.level > tmp_level:
                            tmp_level = tmp_icon.level

                level = tmp_level + 1
                if level > self.maxlevel:
                    self.maxlevel = level
                icon = Icon(idx, imgpath, (left, top), level, over)
                group.append(icon)
                self.group = group
            if not self.keep_use_tool:
                self.used_shuffle = True

    def tool_put(self):
        """ 放回一张图标到原位置 """
        if not self.used_put:
            icon: Icon = self.get_icon_by_idx_in_selected(self.last_selected_icon_idx)
            if icon is None:
                icon = self.selected[-1]
            if icon is not None:
                self.selected.remove(icon)

                # 考虑放回之后的覆盖关系
                tmp_level = 0
                over = []
                for tmp_icon in self.group:
                    if box_iou_xywh((icon.rect.left, icon.rect.top, SIZE, SIZE), tmp_icon.rect) > (1 / 8):
                        tmp_icon.under.append(icon.idx)
                        over.append(tmp_icon.idx)
                        if tmp_icon.level > tmp_level:
                            tmp_level = tmp_icon.level
                level = tmp_level + 1
                if level > self.maxlevel:
                    self.maxlevel = level
                icon.level = level
                icon.over = over
                self.group.append(icon)

                self.remaining += 1
                if not self.keep_use_tool:
                    self.used_put = True

    def tool_withdraw(self):
        """ 抽出三张 放到下面 """
        if not self.used_withdraw and len(self.selected) > 2:
            for i in range(len(self.selected) - 1, max(-1, len(self.selected) - 1 - 3), -1):
                icon = self.selected[i]
                self.selected.remove(icon)
                # 考虑放回之后的覆盖关系
                tmp_level = 0
                over = []
                for tmp_icon in self.group:
                    if box_iou_xywh((icon.rect.left, icon.rect.top, SIZE, SIZE), tmp_icon.rect) > (1 / 8):
                        tmp_icon.under.append(icon.idx)
                        over.append(tmp_icon.idx)
                        if tmp_icon.level > tmp_level:
                            tmp_level = tmp_icon.level
                level = tmp_level + 1
                if level > self.maxlevel:
                    self.maxlevel = level
                icon.level = level
                icon.over = over
                self.group.append(icon)
            self.remaining += 3
            if not self.keep_use_tool:
                self.used_withdraw = True

    def get_icon_by_idx_in_group(self, idx):
        for icon in self.group:
            if icon.idx == idx:
                return icon
        return None

    def get_icon_by_idx_in_selected(self, idx):
        for icon in self.selected:
            if icon.idx == idx:
                return icon
        return None

    def get_icon_by_position(self, position):
        """ 根据鼠标点击的位置获取图标 已经判断了是否可点击（是否还有覆盖该图标的其它图标 没有则可以点击） """
        x, y = position
        for tmp_icon in self.group:
            left, top = tmp_icon.rect.left, tmp_icon.rect.top
            if left < x < left + SIZE and top < y < top + SIZE:
                if len(tmp_icon.under) == 0:
                    self._click_sound()
                    return tmp_icon
                else:
                    usable = True
                    for tmp_idx in tmp_icon.under:
                        if self.get_icon_by_idx_in_group(tmp_idx) is not None:
                            # 还有图标覆盖点击的图标
                            usable = False
                            break
                    if usable:
                        self._click_sound()
                        return tmp_icon
        return None

    def failed(self):
        """ 游戏失败 """
        self.failed_game = True
        self.finished = True

    def success(self):
        """ 游戏成功提示 """
        self.success_game = True
        self.finished = True

    def start(self):
        """ 开始游戏"""
        self.window = pygame.display.set_mode((WIDTH, HEIGHT))
        pygame.mixer.music.play(-1)
        while True:
            self.window.fill((206, 252, 153))  # 窗口背景颜色
            self._draw_selection((MINX, SELECTIONY))  # 绘制七个队列框
            self._draw_remaining_icon_amount()  # 剩余图标数量显示

            # 绘制图标
            for tmp_icon in self.group:
                self.window.blit(tmp_icon.image, tmp_icon.rect)
                # 判断该图标是否可点击 如果不可点击 则加上阴影
                reachable = True  # 默认在顶层
                for tmp_idx in tmp_icon.under:
                    if self.get_icon_by_idx_in_group(tmp_idx) is not None:
                        reachable = False  # unreachable
                        break
                if not reachable:
                    self._draw_transparency_block((tmp_icon.rect.left, tmp_icon.rect.top))  # 给图标添加阴影

            # 绘制候选框中的图标
            for i in range(len(self.selected)):
                self.window.blit(self.selected[i].image, (MINX + i * SIZE, SELECTIONY, SIZE, SIZE))

            self._draw_tools()

            if self.finished:
                self._draw_result()
            for event in pygame.event.get():
                if event.type == pygame.QUIT or (event.type == pygame.KEYUP and event.key == pygame.K_ESCAPE):
                    # 关闭窗口或者按‘Q’键退出游戏
                    pygame.quit()
                    sys.exit()
                if event.type == pygame.KEYUP and event.key == pygame.K_u:
                    self.keep_use_tool = True
                    self.used_shuffle = False
                    self.used_put = False
                    self.used_withdraw = False
                if event.type == pygame.KEYUP and event.key == pygame.K_r:
                    self.reset()
                elif event.type == pygame.MOUSEBUTTONUP:
                    if not self.finished:  # 游戏未结束
                        # 鼠标点击
                        position = pygame.mouse.get_pos()
                        icon_by_position: Icon = self.get_icon_by_position(position)
                        if icon_by_position is not None:
                            # 点击有效 修改该图标的位置
                            self.group.remove(icon_by_position)
                            # 移除原有的覆盖关系
                            for over_icon_idx in icon_by_position.over:
                                over_icon = self.get_icon_by_idx_in_group(over_icon_idx)
                                if over_icon is not None:
                                    over_icon.under.remove(icon_by_position.idx)
                            icon_by_position.over.clear()
                            # 移除原有的被覆盖关系
                            for under_icon_idx in icon_by_position.under:
                                under_icon = self.get_icon_by_idx_in_group(under_icon_idx)
                                if under_icon is not None:
                                    under_icon.over.remove(icon_by_position.idx)
                            icon_by_position.under.clear()
                            # 按照类别进行插入
                            inserted = False  # 是否已插入
                            for i in range(len(self.selected)):
                                tmp_icon = self.selected[i]
                                if tmp_icon.type == icon_by_position.type:
                                    self.selected.insert(i, icon_by_position)
                                    inserted = True
                                    break
                            if not inserted:
                                self.selected.append(icon_by_position)

                            self.last_selected_icon_idx = icon_by_position.idx
                            self.remaining -= 1
                            # 去除匹配成功的图标
                            self._match()
                            # 判断选择框是否已满
                            if len(self.selected) == 7:
                                self.failed()  # 游戏结束
                            if self.remaining == 0:
                                self.success()  # 游戏通关
                        # 点击道具
                        tool_by_position = get_tool_by_position(position)
                        if tool_by_position == 'put':
                            self.tool_put()
                        if tool_by_position == 'shuffle':
                            self.tool_shuffle()
                        if tool_by_position == 'withdraw':
                            self.tool_withdraw()

            pygame.display.update()


def make_train_data(num_iter: int, save_dir: str):
    game = Game()
    game.window = pygame.display.set_mode((WIDTH, HEIGHT))
    for i in range(num_iter):
        game.reset()  # 生成新的图片布局
        game.window.fill((206, 252, 153))  # 窗口背景颜色

        # 绘制图标
        for tmp_icon in game.group:
            game.window.blit(tmp_icon.image, tmp_icon.rect)
            # 判断该图标是否可点击 如果不可点击 则加上阴影
            # reachable = True  # 默认在顶层
            # for tmp_idx in tmp_icon.under:
            #     if game.get_icon_by_idx_in_group(tmp_idx) is not None:
            #         reachable = False  # unreachable
            #         break
            # if not reachable:
            #     game._draw_transparency_block((tmp_icon.rect.left, tmp_icon.rect.top))  # 给图标添加阴影

        pygame.image.save(game.window, os.path.join(save_dir, str(i) + '-screenshot.jpg'))


def main():
    # game = Game()
    # game.start()
    make_train_data(100, 'train')


if __name__ == '__main__':
    main()
